import os
from modules import base
class Mem(base.Base):
    def __init__(self,host=None):
        self.command=r'dmidecode -t memory | grep -v "Size: No Module Installed"'
        self.host=host
        self.res_list=[]
        super(Mem,self).__init__()

    def windows(self):
        res=os.popen('wmic MEMORYCHIP get Speed,SerialNumber,Capacity').read().splitlines()
        del res[0]
        for i in res:
            if len(i) >1:
                clean=i.split(' ')
                while '' in clean:
                    clean.remove('')
                tmp_dict={}
                size='%sMB'%str(int(clean[0])/1024/1024)
                sn=clean[1]
                speed=clean[2]
                tmp_dict['size']=size
                tmp_dict['speed']=speed
                tmp_dict['sn']=sn
                self.res_list.append(tmp_dict)
        base.data_dict['mem_info']=self.res_list
    def linux(self):
        if self.host==None: #如果=None,则证明是agent模式
            res=os.popen(self.command).read()
            self.parse(res)
            base.data_dict['mem_info']=self.res_list

        else: #传入了host_dict，证明是ssh模式
            res=self.ssh(self.host,self.command).decode(encoding='utf-8')
            self.parse(res)
            base.data_dict[self.host]['mem_info']=self.res_list
    def parse(self,data):
        res=data.replace('\n\t','').replace('\n','').split('Device')
        import re
        for i in res:
            if re.match('Array',i):
                tmp_dict={}
                size=re.search('bitsSize.*MB',i)
                if size is not None:
                    tmp_dict['size']=size.group().split(':')[1]
                    speed=re.search('Speed.*?MHz',i)
                    if speed is None:
                        tmp_dict['speed']='Unknown' #有时候内存速度无法获取到，会报错！
                    else:
                        tmp_dict['speed']=speed.group().split(':')[1]

                    sn=re.search('Serial Number.*?Asset',i)
                    if sn is None:
                        tmp_dict['sn']='Unknown'
                    else:
                        tmp_dict['sn']=sn.group().split(':')[1].strip('Asset')
                    self.res_list.append(tmp_dict)



